package com.citusdata.elven.datastore;

import com.citusdata.elven.datastore.DatastoreKey;
import com.citusdata.elven.datastore.DatastoreKeyComparator;
import com.citusdata.elven.datastore.DatastoreValue;
import com.citusdata.elven.datastore.protobuf.Datastore.*;
import com.citusdata.elven.util.ConfigManager;
import com.google.protobuf.ByteString;
import com.google.protobuf.RpcCallback;
import com.google.protobuf.RpcController;
import com.sleepycat.je.Cursor;
import com.sleepycat.je.Database;
import com.sleepycat.je.DatabaseConfig;
import com.sleepycat.je.DatabaseEntry;
import com.sleepycat.je.DatabaseException;
import com.sleepycat.je.Environment;
import com.sleepycat.je.EnvironmentConfig;
import com.sleepycat.je.LockMode;
import com.sleepycat.je.OperationStatus;
import java.io.File;
import java.util.ArrayList;
import java.util.List;
import org.apache.commons.configuration.Configuration;
import org.apache.log4j.Logger;


// Simple datastore implementation backed by the BDB JE engine. This class only
// implements get and put operations, and for now omits deadlock detection and
// data indexing. For deadlock detection, since write queries only touch one
// record at a time, running into deadlocks is unlikely.
// For data indexing, future use cases may require data access by secondary keys
// or (only) time intervals, and creating indices on these dimensions may be
// necessary. For example, analyzing clickstream data for counting the number of
// unique visitors for a given day requires an index over event timestamps.

public class BdbDatastore extends DatastoreInterface {
  private static final Logger logger = Logger.getLogger(BdbDatastore.class);
  private static final Configuration config = ConfigManager.getConfig();
  private static final String DB_NAME = "datastore";

  private final DatastoreKeyComparator keyComparator;
  private final String  directoryName;
  private final boolean transactional;

  private Environment environment;
  private Database    database;


  public BdbDatastore(String directoryName) {
    this.directoryName = directoryName;
    this.keyComparator = new DatastoreKeyComparator();
    this.transactional = config.getBoolean("BdbDatastore.Transactional", true);

    // Initialize and open the BDB JE Engine
    try {
      // Increase lock timeout to 1 second to reduce the amount of lock
      // expirations that may occur.  
      long lockTimeout = 1000*1000;  // 1M microseconds

      // By default, BDB JE Environment Configuration enables checksums, and
      // disables duplicate primary keys. We rely on these defaults; checksums
      // are used to detect data corruption issues, and no-duplication allows
      // for secondary indices on data.

      EnvironmentConfig environmentConfig = new EnvironmentConfig();
      environmentConfig.setAllowCreate(true);
      environmentConfig.setLockTimeout(lockTimeout);
      if (transactional) {
        environmentConfig.setTransactional(true);
        environmentConfig.setTxnWriteNoSync(true);  // Use OS file system buffers
      }
      environment = new Environment(new File(directoryName), environmentConfig);

      // Set custom sort order (comparator) for keys when opening the database
      DatabaseConfig databaseConfig = new DatabaseConfig();
      databaseConfig.setAllowCreate(true);
      databaseConfig.setBtreeComparator(DatastoreKeyComparator.class);
      if (transactional) {
        databaseConfig.setTransactional(true);
      } else {
        databaseConfig.setDeferredWrite(true);
      }
      database = environment.openDatabase(null, DB_NAME, databaseConfig);

      logger.info("Opened database environment at path: " + directoryName);

    } catch (DatabaseException ex) {
      logger.error("Database exception for environment: " + directoryName, ex);
      close();

      // The caller cannot reasonably recover from this failure.
      throw new IllegalStateException("BdbDatastore initialization failed");
    }
  }


  public void close() {
    try {
      if (database != null) {
        database.close();
      }
      if (environment != null) {
        environment.close();
      }
    } catch (DatabaseException ex) {
      logger.info("Database exception during close: " + ex);
    }

    logger.info("Closed environment and databases");
  }


  public boolean put(DatastoreKey key, DatastoreValue value) {
    boolean success = false;

    try {
      DatabaseEntry bdbKey   = new DatabaseEntry(key.getBytes());
      DatabaseEntry bdbValue = new DatabaseEntry(value.getBytes());

      // (1) Since the database was opened as transactional, this write is
      // performed using auto commit. 
      // (2) Since the database does not allow duplicates, if an entry with this
      // key already exists in the database, this put will overwrite it.
      OperationStatus putStatus = database.put(null, bdbKey, bdbValue);

      if (putStatus == OperationStatus.SUCCESS) {
        success = true;
      } else {
        logger.info("Database put failed with status: " + putStatus);
      }
    } catch (DatabaseException ex) {
      logger.warn("Database put failed with exception: " + ex);
    }

    return success;
  }


  public List getRange(DatastoreKey startKey, DatastoreKey endKey) {
    ArrayList entries = new ArrayList();
    Cursor cursor = null;

    try {
      // Create BDB DatabaseEntry objects to iterate over entries
      DatabaseEntry iterKey   = new DatabaseEntry(startKey.getBytes());
      DatabaseEntry iterValue = new DatabaseEntry();

      // Open a cursor, and position it to the first key greater than or equal
      // to the startKey.
      cursor = database.openCursor(null, null);
      OperationStatus status = cursor.getSearchKeyRange(iterKey, iterValue,
                                                        LockMode.DEFAULT);
      while (status == OperationStatus.SUCCESS) {
        // getRange() defines a policy where both endpoints are included in the
        // search; that is, the method returns a closed range.  
        int passedEnd = keyComparator.compare(iterKey.getData(), endKey.getBytes());
        if (passedEnd > 0) {
          break;
        }
        
        DatastoreKey currentKey = new DatastoreKey(iterKey.getData());
        DatastoreValue currentValue = new DatastoreValue(iterValue.getData()); 
        DatastoreEntry currentEntry = new DatastoreEntry(currentKey, currentValue);
        entries.add(currentEntry);

        status = cursor.getNext(iterKey, iterValue, LockMode.DEFAULT);
      }

    } catch (DatabaseException ex) {
      logger.warn("Database range get failed with exception: " + ex);
    } finally {
      try {
        if (cursor != null) {
          cursor.close();
        }
      } catch (DatabaseException ignore) { }
    }
    
    return entries;
  }


  public PutResponse put(PutRequest request) {
    DatastoreKey key = new DatastoreKey(request.getPartitionKey(), request.getTimestamp()); 
    DatastoreValue value = new DatastoreValue(request.getData().toByteArray());

    boolean success = put(key, value);

    PutResponse.Builder putResponse = PutResponse.newBuilder();
    putResponse.setSuccess(success);
    return (putResponse.build());
  }


  public GetRangeResponse getRange(GetRangeRequest request) {
    String partitionKey = request.getPartitionKey();
    DatastoreKey startKey = new DatastoreKey(partitionKey, request.getStartTime());
    DatastoreKey endKey   = new DatastoreKey(partitionKey, request.getEndTime());
    
    GetRangeResponse.Builder rangeResponse = GetRangeResponse.newBuilder();

    List entries = getRange(startKey, endKey);

    for (int i = 0; i < entries.size(); ++i) {
      DatastoreEntry currentEntry = (DatastoreEntry) entries.get(i);

      DatastoreKey currentKey = currentEntry.getDatastoreKey();
      DatastoreValue currentValue = currentEntry.getDatastoreValue();

      PutRequest.Builder putRequest = PutRequest.newBuilder();
      putRequest.setPartitionKey(currentKey.getPartitionKey());
      putRequest.setTimestamp(currentKey.getTimestamp());
      putRequest.setData(ByteString.copyFrom(currentValue.getBytes()));
      
      rangeResponse.addPutRequests(putRequest.build());
    }

    return (rangeResponse.build());
  }


  // Interface methods for remote service calls
  public void put(RpcController controller, PutRequest request, 
                  RpcCallback done) {
    PutResponse response = put(request);
    done.run(response);
  }

  public void getRange(RpcController controller, GetRangeRequest request, 
                       RpcCallback done) {
    GetRangeResponse response = getRange(request);
    done.run(response);
  }
}
